"""
Wan2.1-T2V-14B 模型加载器
使用 MindIE 在 NPU 上进行推理
参考: https://www.hiascend.com/document/detail/zh/mindie/22RC1/quickstart/mindiesd_quickstart_0004.html
"""
import os
import sys
import logging
import yaml
from typing import Optional, Dict

logger = logging.getLogger(__name__)


class Wan2ModelLoader:
    """Wan2.1-T2V模型加载器"""
    
    def __init__(self, config_path: Optional[str] = None):
        # 加载配置
        if config_path is None:
            config_path = os.path.join(os.path.dirname(__file__), "../config/config.yaml")
        
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        model_config = self.config.get('model', {})
        self.model_id = model_config.get('model_id', 'Wan2.1-T2V-14B')
        self.model_path = model_config.get('model_path', '/data2/models/Wan2.1-T2V-14B')
        self.wan2_code_path = model_config.get('wan2_code_path', '/workspace/Wan2.1')
        
        # 推理参数
        self.task = model_config.get('task', 't2v-14B')
        self.size = model_config.get('size', '1280*720')
        self.sample_steps = model_config.get('sample_steps', 50)
        
        # 分布式配置（torchrun）- 2卡
        self.nproc_per_node = model_config.get('nproc_per_node', 2)
        self.ulysses_size = model_config.get('ulysses_size', 2)
        self.dit_fsdp = model_config.get('dit_fsdp', True)
        self.t5_fsdp = model_config.get('t5_fsdp', True)
        self.vae_parallel = model_config.get('vae_parallel', True)
        
        # Attention Cache配置
        self.use_attention_cache = model_config.get('use_attention_cache', True)
        self.start_step = model_config.get('start_step', 20)
        self.attention_cache_interval = model_config.get('attention_cache_interval', 2)
        self.end_step = model_config.get('end_step', 47)
        
        # 设备配置（2卡）
        self.device_ids = model_config.get('device_ids', [1, 2])
        
        self.is_loaded = False
        
    def load_model(self):
        """加载Wan2.1-T2V模型（验证路径和环境）"""
        try:
            logger.info(f"开始加载Wan2.1-T2V模型: {self.model_id}")
            logger.info(f"模型路径: {self.model_path}")
            logger.info(f"Wan2代码路径: {self.wan2_code_path}")
            
            # 验证模型路径
            if not os.path.exists(self.model_path):
                raise FileNotFoundError(f"模型路径不存在: {self.model_path}")
            
            # 验证Wan2代码路径
            if not os.path.exists(self.wan2_code_path):
                raise FileNotFoundError(f"Wan2代码路径不存在: {self.wan2_code_path}")
            
            # 验证generate.py存在
            generate_script = os.path.join(self.wan2_code_path, 'generate.py')
            if not os.path.exists(generate_script):
                raise FileNotFoundError(f"generate.py不存在: {generate_script}")
            
            # 添加Wan2代码路径到Python路径
            if self.wan2_code_path not in sys.path:
                sys.path.insert(0, self.wan2_code_path)
            
            logger.info(f"✓ 模型路径验证通过")
            logger.info(f"✓ Wan2代码路径验证通过")
            
            # 打印推理配置
            logger.info("=" * 60)
            logger.info("推理配置:")
            logger.info(f"  任务类型: {self.task}")
            logger.info(f"  分辨率: {self.size}")
            logger.info(f"  采样步数: {self.sample_steps}")
            logger.info(f"  NPU卡数: {self.nproc_per_node}")
            logger.info(f"  Ulysses Size: {self.ulysses_size}")
            logger.info(f"  DiT FSDP: {self.dit_fsdp}")
            logger.info(f"  T5 FSDP: {self.t5_fsdp}")
            logger.info(f"  VAE Parallel: {self.vae_parallel}")
            logger.info(f"  Attention Cache: {self.use_attention_cache}")
            logger.info(f"  设备ID: {self.device_ids}")
            logger.info("=" * 60)
            
            self.is_loaded = True
            logger.info("Wan2.1-T2V模型加载完成")
            return True
            
        except Exception as e:
            logger.error(f"模型加载失败: {str(e)}", exc_info=True)
            raise
    
    def get_config(self) -> Dict:
        """获取配置"""
        return self.config
    
    def get_model_path(self) -> str:
        """获取模型路径"""
        return self.model_path
    
    def get_wan2_code_path(self) -> str:
        """获取Wan2代码路径"""
        return self.wan2_code_path
    
    def get_inference_params(self) -> Dict:
        """获取推理参数"""
        return {
            'task': self.task,
            'size': self.size,
            'sample_steps': self.sample_steps,
            'nproc_per_node': self.nproc_per_node,
            'ulysses_size': self.ulysses_size,
            'dit_fsdp': self.dit_fsdp,
            't5_fsdp': self.t5_fsdp,
            'vae_parallel': self.vae_parallel,
            'use_attention_cache': self.use_attention_cache,
            'start_step': self.start_step,
            'attention_cache_interval': self.attention_cache_interval,
            'end_step': self.end_step,
            'device_ids': self.device_ids
        }
